from transformers import AutoTokenizer, DataCollatorWithPadding,DataCollatorForTokenClassification,DefaultDataCollator
from datasets import Dataset
from torch.utils.data import DataLoader
import os

'''
DataCollatorForTokenClassification  会自动为 labels标签进行 填充
DataCollatorWithPadding   不会为了 labels 标签进行填充



'''
if os.getlogin() == 'caofei':
    # data_folder = r'C:\Users\caofei\Desktop\desktop link\torch1\hgface\token_classification\ner_data'
    model_folder = r'D:\models\chinese-macbert-base'
else:
    model_folder = r'D:\Models\chinese-macbert-base'

# 1. 准备示例数据（不同长度的文本）
data = [
    {"text": "I love natural language processing", "label": 1},  # 短文本
    {"text": "Hugging Face's transformers library is powerful for building NLP models. It supports many state-of-the-art models like BERT, GPT, and T5.", "label": 0},  # 长文本
    {"text": "DataCollatorWithPadding helps pad sequences to the same length in a batch", "label": 1},  # 中等长度
    {"text": "Hello world", "label": 0}  # 极短文本
]

data = [
    {
        "text": "Alice lives in Paris",
        "labels": [1, 0, 0, 2]  # Alice(人名)=1，lives=0，in=0，Paris(地点)=2
    },
    {
        "text": "Bob works at Google asd",
        "labels": [1, 0, 0, 3,4]  # Bob(人名)=1，works=0，at=0，Google(机构)=3
    }
]


# 转换为datasets格式的Dataset
dataset = Dataset.from_list(data)


# 2. 加载Tokenizer（以BERT为例）
tokenizer = AutoTokenizer.from_pretrained(model_folder)


# 3. 定义预处理函数（仅分词，不提前填充）
def preprocess_function(examples):
    # 只做分词和截断，不设置padding=True（填充交给DataCollatorWithPadding）
    return tokenizer(
        examples["text"],
        truncation=True,  # 超过max_length的部分截断
        max_length=512,    # BERT的默认最大长度
        # padding=True,
        # return_tensors='pt'
    )

# 对数据集应用预处理（得到分词后的结果，如input_ids、attention_mask）
tokenized_dataset = dataset.map(preprocess_function, batched=True,remove_columns=['text'])

print(tokenized_dataset)


for d in tokenized_dataset:
    print(d['input_ids'])

# # 4. 初始化DataCollatorWithPadding（核心：自动填充）
# # 传入tokenizer，它会自动使用模型的pad_token（BERT中是[PAD]，id=0）
# data_collator = DataCollatorWithPadding(tokenizer=tokenizer)
data_collator = DataCollatorForTokenClassification(tokenizer=tokenizer)


# # 5. 将数据集转换为PyTorch格式（方便DataLoader处理）
# tokenized_dataset = tokenized_dataset.set_format(
#     "torch",
#     columns=["input_ids", "attention_mask", "label"]  # 指定需要的列
# )


# # 6. 用DataLoader加载数据（结合collate_fn实现填充）
dataloader = DataLoader(
    tokenized_dataset,
    batch_size=2,  # 每批2个样本
    collate_fn=data_collator  # 用DataCollatorWithPadding处理填充
)



# d = next(iter(dataloader))
# print(d)

# # 7. 查看填充效果
for batch in dataloader:
    print("=== 新批次 ===")
    print("input_ids形状：", batch["input_ids"].shape)  # 所有样本长度相同（填充后）
    print("第一个样本input_ids：", batch["input_ids"][0])  # 包含实际文本+填充的0（[PAD]）
    print("第一个样本attention_mask：", batch["attention_mask"][0])  # 1表示有效文本，0表示填充
    print("第二个样本input_ids：", batch["input_ids"][1])
    print("第二个样本attention_mask：", batch["attention_mask"][1])
    # print("标签：", batch["label"])
    print("\n")